import tkinter as tk
import os
import pandas as pd
from sklearn.neighbors import KNeighborsClassifier
import numpy as np
import csv
from PIL import Image, ImageTk
import logging
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
logging.basicConfig(level=logging.INFO, format='%(asctime)s:%(levelname)s:%(message)s')

class MovieRatingApp:
    def __init__(self, root):
        self.root = root
        root.title("Movie Rating App")
        root.geometry("900x600")

        # 图片加载（占位图逻辑）
        try:
            self.image1 = Image.open(os.path.join(os.path.dirname(__file__), 'post1_relie.jpg'))
        except:
            self.image1 = Image.new('RGB', (100, 300), color='orange')
        self.image1.thumbnail((100, 300))
        self.movie1_image = ImageTk.PhotoImage(self.image1)
        self.movie1_label = tk.Label(root, image=self.movie1_image)
        self.movie1_label.grid(row=0, column=0, padx=10, pady=10)

        try:
            self.image2 = Image.open(os.path.join(os.path.dirname(__file__), "post2_family.jpg"))
        except:
            self.image2 = Image.new('RGB', (100, 300), color='blue')
        self.image2.thumbnail((100, 300))
        self.movie2_image = ImageTk.PhotoImage(self.image2)
        self.movie2_label = tk.Label(root, image=self.movie2_image)
        self.movie2_label.grid(row=0, column=1, padx=10, pady=10)

        # 评分下拉菜单
        self.movie1_rating = tk.StringVar(root, value="1")
        tk.OptionMenu(root, self.movie1_rating, "1", "2", "3", "4", "5").grid(row=1, column=0, padx=10, pady=10)
        self.movie2_rating = tk.StringVar(root, value="1")
        tk.OptionMenu(root, self.movie2_rating, "1", "2", "3", "4", "5").grid(row=1, column=1, padx=10, pady=10)

        # 电影类型单选按钮
        self.preference = tk.IntVar()
        tk.Radiobutton(root, text="动作片", variable=self.preference, value=0).grid(row=2, column=0, padx=10, pady=10)
        tk.Radiobutton(root, text="喜剧片", variable=self.preference, value=1).grid(row=2, column=1, padx=10, pady=10)

        # 功能按钮
        tk.Button(root, text="Confirm", command=self.confirm).grid(row=3, column=0, padx=10, pady=20)
        tk.Button(root, text="Clear", command=self.clear_plot).grid(row=3, column=1, padx=10, pady=20)
        tk.Button(root, text="Predict", command=self.predict_preference).grid(row=4, columnspan=2, padx=10, pady=20)
        self.predict_label = tk.Label(root, text="")
        self.predict_label.grid(row=5, columnspan=2, padx=10, pady=10)

        # 数据文件路径
        self.filename = os.path.join(os.path.dirname(__file__), 'ratings.csv')
        if not os.path.isfile(self.filename):
            with open(self.filename, 'w', newline='') as file:
                writer = csv.writer(file)
                writer.writerow(["Movie1", "Movie2", "Preference"])

        # 图表相关变量
        self.plot_frame = tk.Frame(root)
        self.plot_frame.grid(row=0, column=2, rowspan=4, padx=10, pady=10)
        self.fig = None  # 保存当前图表
        self.ax = None   # 保存当前坐标轴
        self.canvas = None  # 保存当前画布
        self.plot_ratings()

    def confirm(self):
        m1_rating = self.movie1_rating.get()
        m2_rating = self.movie2_rating.get()
        preference = self.preference.get()
        with open(self.filename, 'a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([m1_rating, m2_rating, preference])
        self.plot_ratings()

    def clear_plot(self):
        with open(self.filename, 'w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(["Movie1", "Movie2", "Preference"])
        self.predict_label.config(text="")
        self.plot_ratings()

    def plot_ratings(self):
        # 清除现有图表
        for widget in self.plot_frame.winfo_children():
            widget.destroy()
        
        # 创建新图表和坐标轴
        self.fig, self.ax = plt.subplots(figsize=(5, 4))
        
        df = pd.read_csv(self.filename)
        
        if len(df) <= 1:  # 只有表头，无数据
            self.ax.set_xlabel('Movie A Rating')
            self.ax.set_ylabel('Movie B Rating')
            self.ax.set_xlim([0, 6])
            self.ax.set_ylim([0, 6])
            self.ax.grid(True)
            self.ax.legend(['No data yet'])
        else:  # 有数据时
            action_movies = df[df['Preference'] == 0]
            comedy_movies = df[df['Preference'] == 1]
            
            # 绘制历史数据点
            self.ax.scatter(action_movies['Movie1'], action_movies['Movie2'], 
                           color='orange', label='Action')
            self.ax.scatter(comedy_movies['Movie1'], comedy_movies['Movie2'], 
                           color='blue', label='Comedy')
            
            # 添加坐标文本
            for _, row in action_movies.iterrows():
                self.ax.text(row['Movie1'], row['Movie2'] + 0.1, 
                           f'({row["Movie1"]}, {row["Movie2"]})', fontsize=8)
            for _, row in comedy_movies.iterrows():
                self.ax.text(row['Movie1'], row['Movie2'] + 0.1, 
                           f'({row["Movie1"]}, {row["Movie2"]})', fontsize=8)
            
            # 设置坐标轴
            self.ax.set_xlabel('Movie A Rating')
            self.ax.set_ylabel('Movie B Rating')
            self.ax.set_xlim([0, 6])
            self.ax.set_ylim([0, 6])
            self.ax.grid(True)
            self.ax.legend()
        
        # 创建并显示画布
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.plot_frame)
        self.canvas.draw()
        self.canvas.get_tk_widget().pack()

    def predict_preference(self):
        df = pd.read_csv(self.filename)
        
        # 检查是否有足够数据
        if len(df) <= 1:
            self.predict_label.config(text="请先添加一些评分数据再进行预测")
            return
            
        # 准备数据
        X = df[['Movie1', 'Movie2']].values.astype(int)
        y = df['Preference'].values
        
        # 训练KNN模型
        knn = KNeighborsClassifier(n_neighbors=1)
        knn.fit(X, y)
        
        # 获取新用户评分
        m1_rating = int(self.movie1_rating.get())
        m2_rating = int(self.movie2_rating.get())
        new_point = np.array([[m1_rating, m2_rating]])
        
        # 预测
        prediction = knn.predict(new_point)
        self.predict_label.config(text=f"给用户推荐的电影类型: {'动作片' if prediction[0] == 0 else '喜剧片'}")
        
        # 获取最近邻
        distances, indices = knn.kneighbors(new_point)
        nearest_neighbor_index = indices[0][0]
        nearest_neighbor = X[nearest_neighbor_index]
        
        # 在现有图表上直接绘制新点和连接线（关键修复）
        # 绘制新用户点
        self.ax.scatter(new_point[0, 0], new_point[0, 1], 
                       color='red', marker='x', s=100, label='New User')
        
        # 绘制连接线（使用明显的样式）
        self.ax.plot([new_point[0, 0], nearest_neighbor[0]],
                    [new_point[0, 1], nearest_neighbor[1]],
                    'r--', linewidth=2, label='Nearest Neighbor')
        
        # 更新图例
        self.ax.legend()
        
        # 强制刷新画布
        self.canvas.draw()

root = tk.Tk()
app = MovieRatingApp(root)
root.mainloop()
